{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 2 - Build RAG-powered Q&A Application with **RetrieveAndGenerate API**\n",
    "\n",
    "----\n",
    "\n",
    "This notebook provides sample code and step-by-step instructions for building a fully-managed question-answering (Q&A) application using a **RetrieveAndGenerate API** of Amazon Bedrock Knowledge Bases.\n",
    "\n",
    "----\n",
    "\n",
    "### Introduction\n",
    "\n",
    "In the previous notebook, we demonstrated how to create a Knowledge Base in Amazon Bedrock — including setting up an S3 data source, configuring an Amazon OpenSearch Serverless (AOSS) vector index, and ingesting documents for retrieval-augmented generation (RAG).\n",
    "\n",
    "In this notebook, we take the next step: building a Q&A application that can query that Knowledge Base using the `RetrieveAndGenerate` API. This API allows you to retrieve the most relevant content from your Knowledge Base based on a user’s query and automatically pass that information to a foundation model (FM) to generate a grounded, context-aware response.\n",
    "\n",
    "This is a classic example of the RAG pattern — where external data is dynamically retrieved at query time and incorporated into the model’s prompt to improve relevance, accuracy, and transparency. In this solution, retrieved knowledge base content comes with source attribution, helping end users understand the origin of the response and minimizing the risk of model hallucinations.\n",
    "\n",
    "![BKB illustration](./images/retrieve_and_generate_api.png)\n",
    "\n",
    "### Pre-requisites\n",
    "\n",
    "In order to run this notebook, you should have successfully completed the previous notebook lab:\n",
    "- [1_create-kb-and-ingest-documents.ipynb](./1\\_create-kb-and-ingest-documents.ipynb).\n",
    "\n",
    "Also, please make sure that you have enabled the following model access in _Amazon Bedrock Console_:\n",
    "- `Amazon Nova Micro`\n",
    "- `Amazon Titan Text Embeddings V2`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "### 1.1 Import the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard library imports\n",
    "import os\n",
    "import sys\n",
    "import json\n",
    "import time\n",
    "\n",
    "# Third-party imports\n",
    "import boto3\n",
    "from botocore.client import Config\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# Local imports\n",
    "import utility\n",
    "\n",
    "# Print SDK versions\n",
    "print(f\"Python version: {sys.version.split()[0]}\")\n",
    "print(f\"Boto3 SDK version: {boto3.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Initial setup for clients and global variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r bedrock_kb_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create boto3 session and set AWS region\n",
    "boto_session = boto3.Session()\n",
    "aws_region = boto_session.region_name\n",
    "\n",
    "# Create boto3 clients for Bedrock\n",
    "bedrock_config = Config(connect_timeout=120, read_timeout=120, retries={'max_attempts': 0})\n",
    "bedrock_agent_client = boto3.client('bedrock-agent-runtime', config=bedrock_config)\n",
    "\n",
    "# Set the Bedrock model to use for text generation\n",
    "model_id = 'amazon.nova-micro-v1:0'\n",
    "model_arn = f'arn:aws:bedrock:{aws_region}::foundation-model/{model_id}'\n",
    "\n",
    "# Print configurations\n",
    "print(\"AWS Region:\", aws_region)\n",
    "print(\"Bedrock Knowledge Base ID:\", bedrock_kb_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Fully-managed RAG with **RetreiveAndGenerate API**\n",
    "\n",
    "The `RetrieveAndGenerate` API provides a fully managed way to implement the Retrieval-Augmented Generation (RAG) pattern with Amazon Bedrock Knowledge Bases.\n",
    "\n",
    "When a user submits a query, the API automatically converts the query into vector embeddings, performs a similarity search against the Knowledge Base, and retrieves the most relevant document chunks. These search results are then injected into the foundation model's prompt as additional context, enabling the model to generate more accurate and grounded responses.\n",
    "\n",
    "For multi-turn conversations, Knowledge Bases also maintain short-term conversational memory — allowing the API to return more contextually relevant answers across a dialogue.\n",
    "\n",
    "The output of the `RetrieveAndGenerate` API includes:\n",
    "\n",
    "- The **generated response** from the foundation model\n",
    "\n",
    "- **Source attribution** metadata for the retrieved content\n",
    "\n",
    "- The **actual retrieved text chunks** from the Knowledge Base\n",
    "\n",
    "This makes it easy to build RAG-powered applications with trusted, explainable answers — without having to manage retrieval pipelines or prompt construction yourself.\n",
    "\n",
    "### 2.1 Retrieve and Generate Example\n",
    "\n",
    "Let’s now see the `RetrieveAndGenerate` API in action and showcase a fully managed RAG workflow in Amazon Bedrock.\n",
    "\n",
    "In this example, we’ll use the Knowledge Base built in the previous lab — containing Amazon Shareholder Letters — to demonstrate how the API retrieves relevant information and generates a grounded response to a user query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_query = \"How does Amazon use technology to better serve its customers?\"\n",
    "\n",
    "response = bedrock_agent_client.retrieve_and_generate(\n",
    "    input={\n",
    "        'text': user_query\n",
    "    },\n",
    "    retrieveAndGenerateConfiguration={\n",
    "        'type': 'KNOWLEDGE_BASE',\n",
    "        'knowledgeBaseConfiguration': {\n",
    "            'knowledgeBaseId': bedrock_kb_id,\n",
    "            'modelArn': model_arn\n",
    "        }\n",
    "    }\n",
    ")\n",
    "\n",
    "print(\"Final reply:\\n\", response['output']['text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Understanding Citations\n",
    "\n",
    "Citations play a critical role in retrieval-augmented generation (RAG) systems by helping users verify the accuracy of a response and providing transparency into the source of information. Let's now look at the `citations` past of the Knowledge Base response:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Citations:\\n\", json.dumps(response[\"citations\"], indent=2, default=str))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the response includes a `generatedResponsePart` field, which contains the natural language answer generated by the model. Each `generatedResponsePart` is paired with `retrievedReferences`, which lists the specific pieces of content from the knowledge base that were used to ground that part of the response. These references include the original source text (`content.text`), as well as metadata like the source URI and page number, so users can easily trace information back to its original document. This structure ensures that answers are both helpful and verifiable, allowing users to explore the source material directly and build trust in the response."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Conclusions and Next Steps\n",
    "\n",
    "In this notebook, we built a fully-managed RAG-powered Q&A application using the `RetrieveAndGenerate` API from Amazon Bedrock Knowledge Bases.\n",
    "\n",
    "We demonstrated how this API simplifies the RAG workflow by automatically retrieving relevant content from a knowledge base and generating grounded, context-aware responses using a foundation model. The responses also include source references, allowing users to easily trace answers back to the original documents.\n",
    "\n",
    "This approach enables you to quickly build reliable, transparent Q&A solutions without managing the complexity of prompt engineering or retrieval logic manually.\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "If you are looking for more flexibility and control over your RAG workflow, Amazon Bedrock Knowledge Bases also provides a `Retrieve` API. This API allows you to perform semantic and/or keyword search over your knowledge base and retrieve the most relevant document chunks, which you can then use to build custom prompts or workflows tailored to your application needs.\n",
    "\n",
    "To explore this approach, check out the next notebook:\n",
    "\n",
    "&nbsp; **NEXT ▶** [3_customized-rag-with-retrieve-api](./3\\_customized-rag-with-retrieve-api.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
